"""
PATH Solver Interface

实现Julia PATH求解器的Python接口，参考DGSQP中的实现
"""

import numpy as np
import torch
from typing import Dict, Any, Optional, Tuple, Callable
import warnings

try:
    from julia.api import Julia
    from julia import Main
    JULIA_AVAILABLE = True
except ImportError:
    JULIA_AVAILABLE = False
    warnings.warn("Julia not available. PATH solver will not work.")


class PATHSolver:
    """
    PATH求解器的Python包装器

    基于DGSQP仓库中的实现，提供Julia PATH求解器的Python接口
    
    性能优化：
    - Julia函数定义只执行一次（初始化时）
    - 边界条件和参数只设置一次
    - 使用全局变量传递函数引用，避免PyJulia闭包转换开销
    """

    def __init__(self, tolerance: float = 1e-6, verbose: bool = False):
        """
        初始化PATH求解器

        Args:
            tolerance: 收敛容差
            verbose: 是否显示详细输出
        """
        self.tolerance = tolerance
        self.verbose = verbose
        self.initialized = False
        self.julia_functions_defined = False  # 🆕 标记Julia函数是否已定义
        
        # 🆕 缓存边界条件，避免重复传递
        self._cached_bounds = None  # (lb, ub, n)
        self._cached_problem_size = None
        
        # 🆕 预分配Julia数组，避免每次创建新对象
        self._julia_z0_cached = None
        self._julia_solution_cached = None
        
        # 🔧 缓存求解参数，避免每次都设置
        self._cached_tolerance = None
        self._cached_solve_params = None  # 缓存solve_cmd字符串模板

        if JULIA_AVAILABLE:
            self._initialize_julia()

    def _initialize_julia(self):
        """初始化Julia环境"""
        try:
            # 初始化Julia
            jl = Julia(compiled_modules=False)

            # 导入必要的包
            Main.eval('using PyCall')
            Main.eval('using PATHSolver')

            # 设置许可证（如果需要）
            license_str = "2830898829&Courtesy&&&USR&45321&5_1_2021&1000&PATH&GEN&31_12_2025&0_0_0&6000&0_0"
            Main.eval(f'PATHSolver.c_api_License_SetString("{license_str}")')

            # 🆕 一次性定义Julia函数框架（避免每次调用都重新定义）
            self._define_julia_functions()

            self.initialized = True
            # print("PATH求解器初始化成功")

        except Exception as e:
            warnings.warn(f"PATH求解器初始化失败: {e}")
            self.initialized = False

    def _define_julia_functions(self):
        """
        🆕 一次性定义Julia函数框架
        
        性能优化关键：
        - Julia函数只定义一次，复用多次
        - 使用全局变量存储Python函数引用（避免每次重新序列化）
        - 减少PyJulia闭包转换开销
        """
        if self.julia_functions_defined:
            return
        
        try:
            # 定义F函数（使用全局的F_py_cached引用）
            F_def = """
            function F_cached(n::Cint, x::Vector{Cdouble}, f::Vector{Cdouble})
                @assert n == length(x)
                f .= F_py_cached(x)
                return Cint(0)
            end
            """
            Main.eval(F_def)

            # 定义J函数（使用全局的J_py_cached引用）
            J_def = """
            function J_cached(
                n::Cint,
                nnz::Cint,
                x::Vector{Cdouble},
                col::Vector{Cint},
                len::Vector{Cint},
                row::Vector{Cint},
                data::Vector{Cdouble},
            )
                @assert n == length(x) == length(col) == length(len)
                @assert nnz == length(row) == length(data)
                j = Array{Float64}(undef, n, n)
                j .= J_py_cached(x)
                i = 1
                for c in 1:n
                    col[c], len[c] = i, 0
                    for r in 1:n
                        row[i], data[i] = r, j[r, c]
                        len[c] += 1
                        i += 1
                    end
                end
                return Cint(0)
            end
            """
            Main.eval(J_def)
            
            self.julia_functions_defined = True
            # 🔧 移除打印，避免I/O开销
            # print("  ✓ Julia函数框架已定义（仅一次）")
            
        except Exception as e:
            warnings.warn(f"定义Julia函数失败: {e}")
            self.julia_functions_defined = False
    
    def clear_cache(self):
        """
        🔧 清理PATH求解器的缓存和状态（关键优化）
        
        ⚠️ 在每个epoch后调用，防止PATH求解器性能下降
        
        清理内容：
        1. Python端的缓存标志
        2. Julia端的全局变量（F_py_cached, J_py_cached等）
        3. 触发完整的Julia垃圾回收（GC.gc(true)）
        """
        if not self.initialized:
            return
        
        try:
            # 🔧 重置Python端的缓存标志
            self._cached_bounds = None
            self._cached_problem_size = None
            self._julia_z0_cached = None
            self._julia_solution_cached = None
            self._cached_tolerance = None
            self._cached_solve_params = None
            
            # 🔧 清理Julia端的缓存变量（关键）
            try:
                Main.eval("""
                    # 清理函数缓存（如果存在）
                    if @isdefined(F_py_cached)
                        global F_py_cached = nothing
                    end
                    if @isdefined(J_py_cached)
                        global J_py_cached = nothing
                    end
                    
                    # 清理数据变量（如果存在且不是函数）
                    if @isdefined(z0) && !isa(z0, Function)
                        global z0 = nothing
                    end
                    if @isdefined(lb) && !isa(lb, Function)
                        global lb = nothing
                    end
                    if @isdefined(ub) && !isa(ub, Function)
                        global ub = nothing
                    end
                    if @isdefined(z) && !isa(z, Function)
                        global z = nothing
                    end
                """)
            except:
                pass  # 忽略清理变量的错误
            
            # 🔧 触发完整的Julia垃圾回收（Windows使用true参数）
            Main.eval('GC.gc(true)')  # true表示强制完整GC
            
            # 🔧 移除打印，避免I/O开销
            # print("  ✓ PATH求解器缓存已清理")
            
        except Exception as e:
            # 静默失败，不影响训练
            pass

    def solve_mcp(
        self,
        F_func: Callable,
        J_func: Callable,
        z0: np.ndarray,
        lb: np.ndarray,
        ub: np.ndarray,
        **kwargs
    ) -> Tuple[np.ndarray, bool, float, str]:
        """
        求解混合互补问题（优化版）
        
        性能优化：
        - 使用缓存的Julia函数（F_cached, J_cached）
        - 只更新全局变量（F_py_cached, J_py_cached）
        - 避免每次调用都重新定义Julia函数
        - 每100次调用触发Julia GC，防止对象累积

        Args:
            F_func: 函数F的Python函数
            J_func: 雅可比矩阵J的Python函数
            z0: 初始点
            lb: 下界
            ub: 上界
            **kwargs: 其他求解参数

        Returns:
            - solution: 解
            - success: 是否成功
            - residual: 残差
            - status: 求解状态
        """
        if not self.initialized:
            raise RuntimeError("PATH求解器未初始化")

        # 🔧 记录调用次数，用于定期GC
        if not hasattr(self, '_solve_count'):
            self._solve_count = 0
        self._solve_count += 1

        try:
            n = len(z0)
            
            # 🔧 关键优化：每100次调用清理一次Julia全局变量，防止累积
            # 🔧 移除打印，避免I/O开销
            if self._solve_count % 100 == 0:
                Main.eval("GC.gc(true)")  # 触发完整Julia垃圾回收
                # print(f"  🔄 第{self._solve_count}次调用：触发Julia GC")
            
            # 🆕 只在边界变化或问题规模变化时才传递边界条件
            bounds_changed = (
                self._cached_bounds is None or 
                self._cached_problem_size != n or
                not np.array_equal(self._cached_bounds[0], lb) or
                not np.array_equal(self._cached_bounds[1], ub)
            )
            
            if bounds_changed:
                Main.lb = lb
                Main.ub = ub
                self._cached_bounds = (lb.copy(), ub.copy())
                self._cached_problem_size = n
                # print(f"  🔄 边界条件已更新 (n={n})")
            # else:
            #     print(f"  ✓ 复用缓存的边界条件 (n={n})")
            
            # 🆕 优化：原地更新z0而非每次创建新Julia数组
            if self._julia_z0_cached is None or len(z0) != n:
                Main.z0 = z0
                self._julia_z0_cached = Main.z0
            else:
                # 原地更新Julia数组（避免重新分配内存）
                try:
                    Main.z0[:] = z0  # Julia语法：原地赋值
                except:
                    # 如果失败，回退到创建新数组
                    Main.z0 = z0
                    self._julia_z0_cached = Main.z0
            
            # 🔧 优化：tolerance只在变化时设置（通常不变）
            if self._cached_tolerance != self.tolerance:
                Main.tol = self.tolerance
                self._cached_tolerance = self.tolerance

            # 🆕 关键优化：只更新全局函数引用，不重新定义Julia函数
            Main.F_py_cached = F_func
            Main.J_py_cached = J_func

            # 🔧 关键优化：缓存求解命令字符串，避免每次都构建
            # 只在参数变化时重新构建
            output = "yes" if self.verbose else "no"
            nms = kwargs.get('nms', 'no')
            major_iter_limit = kwargs.get('major_iteration_limit', 500)
            minor_iter_limit = kwargs.get('minor_iteration_limit', 10000)
            cumulative_iter_limit = kwargs.get('cumulative_iteration_limit', 50000)
            time_limit = kwargs.get('time_limit', 300.0)
            restart_limit = kwargs.get('restart_limit', 100)
            
            # 构建参数元组用于比较
            current_params = (output, nms, major_iter_limit, minor_iter_limit, 
                            cumulative_iter_limit, time_limit, restart_limit, n)
            
            if self._cached_solve_params != current_params:
                # 参数变化了，重新构建命令字符串
                self._solve_cmd = f"""
status, z, info = PATHSolver.solve_mcp(
    F_cached, J_cached, lb, ub, z0,
    nnz={n**2},
    output="{output}",
    convergence_tolerance=tol,
    nms="{nms}",
    crash_nbchange_limit=1000,
    major_iteration_limit={major_iter_limit},
    minor_iteration_limit={minor_iter_limit},
    cumulative_iteration_limit={cumulative_iter_limit},
    time_limit={time_limit},
    restart_limit={restart_limit}
)
success = status == PATHSolver.MCP_Solved
"""
                self._cached_solve_params = current_params
                # print(f"  🔄 重新构建求解命令（参数变化）")
            # else:
            #     print(f"  ✓ 复用缓存的求解命令")

            # 执行求解（使用缓存的命令字符串）
            Main.eval(self._solve_cmd)

            # 🔧 优化：复用numpy数组，避免每次创建新对象
            if self._julia_solution_cached is None or len(self._julia_solution_cached) != n:
                # 第一次或尺寸变化：创建新数组
                solution = np.array(Main.z)
                self._julia_solution_cached = solution
            else:
                # 复用已有数组，原地更新（减少内存分配）
                try:
                    self._julia_solution_cached[:] = Main.z
                    solution = self._julia_solution_cached
                except:
                    # 失败则创建新数组
                    solution = np.array(Main.z)
                    self._julia_solution_cached = solution
            
            # 提取标量结果（这些对象很小，不需要缓存）
            success = bool(Main.success)
            residual = float(Main.info.residual)
            status = str(Main.status)

            return solution, success, residual, status

        except Exception as e:
            warnings.warn(f"PATH求解失败: {e}")
            return z0, False, float('inf'), "error"


class PATHMCPWrapper:
    """
    PATH MCP求解器包装器，兼容PyTorch自动微分

    基于DGSQP中的PATHMCP类实现
    
    性能优化：
    - 复用PATHSolver实例（避免重复初始化Julia）
    - 缓存上次求解结果作为warm start
    - 减少Python→Julia的数据传递开销
    """

    def __init__(
        self,
        tolerance: float = 1e-6,
        verbose: bool = False,
        max_iterations: int = 100000,
        major_iteration_limit: int = 500,
        minor_iteration_limit: int = 10000,
        time_limit: float = 60.0
    ):
        self.path_solver = PATHSolver(tolerance, verbose)
        self.tolerance = tolerance
        self.verbose = verbose
        self.max_iterations = max_iterations
        self.major_iteration_limit = major_iteration_limit
        self.minor_iteration_limit = minor_iteration_limit
        self.time_limit = time_limit
        
        # 🆕 性能优化：缓存上次求解结果用于warm start
        self.last_solution = None
        self.last_problem_size = None
    
    def clear_solver_cache(self):
        """
        🔧 清理求解器缓存，防止warm start信息累积导致性能下降（架构级优化）
        
        在长时间训练中，PATH solver的warm start信息可能累积过多，
        定期清理可以防止求解速度变慢
        
        清理内容：
        1. Python端的warm start缓存
        2. 底层PATHSolver的Julia端缓存（关键）
        3. 触发Julia完整GC
        """
        # Python端缓存
        self.last_solution = None
        self.last_problem_size = None
        self._solve_count = 0

        # 🔧 关键：调用底层PATHSolver的清理方法（清理Julia端）
        self.path_solver.clear_cache()

    def solve(
        self,
        F_func: Callable,
        J_func: Callable,
        z0: np.ndarray,
        lb: Optional[np.ndarray] = None,
        ub: Optional[np.ndarray] = None,
        use_warm_start: bool = False  # 🔧 禁用warm start，避免求解器状态累积
    ) -> Dict[str, Any]:
        """
        求解MCP问题（参考PATHMCP.py的_solve_mcp方法）
        
        性能优化：
        - 支持warm start（使用上次求解结果作为初始点）
        - 复用边界条件（如果问题规模相同）

        Args:
            F_func: 函数F
            J_func: 雅可比矩阵函数J
            z0: 初始点
            lb: 下界
            ub: 上界
            use_warm_start: 是否使用warm start

        Returns:
            求解结果字典
        """
        n = len(z0)

        # 默认边界
        if lb is None:
            lb = -np.inf * np.ones(n)
        if ub is None:
            ub = np.inf * np.ones(n)

        # 🆕 性能优化：使用warm start（如果可用且问题规模相同）
        initial_guess = z0
        if use_warm_start and self.last_solution is not None and self.last_problem_size == n:
            initial_guess = self.last_solution
            # print(f"  🚀 使用warm start（复用上次求解结果）")
        else:
            # print(f"  🔄 冷启动（首次求解或问题规模变化）")
            pass

        # print(f"\nPATH求解器参数:")
        # print(f"  - 变量数: {n}")
        # print(f"  - 容差: {self.tolerance}")
        # print(f"  - Major迭代限制: {self.major_iteration_limit}")
        # print(f"  - Minor迭代限制: {self.minor_iteration_limit}")
        # print(f"  - 时间限制: {self.time_limit}s")

        # 求解（使用优化的初始点）
        solution, success, residual, status = self.path_solver.solve_mcp(
            F_func, J_func, initial_guess, lb, ub, 
            nms='no',
            major_iteration_limit=self.major_iteration_limit,
            minor_iteration_limit=self.minor_iteration_limit,
            time_limit=self.time_limit
        )
        
        # 🆕 缓存求解结果用于下次warm start
        if success:
            self.last_solution = solution.copy()
            self.last_problem_size = n

        # 计算统计量（参考PATHMCP.py的计算方式）
        F_val = F_func(solution)
        
        # 分离出控制变量部分和约束部分
        # F = [∇_u L; -g(u)]，其中前n_u部分是梯度，后面是约束
        # 需要根据实际问题确定分界点
        # 这里假设后1/3是约束相关（拉格朗日乘子）
        n_controls = int(n * 2 / 3)
        
        # 静态性：梯度的范数
        stat = np.linalg.norm(F_val[:n_controls], ord=np.inf)
        
        # 可行性：约束违反（g(u) < 0的部分）
        g_val = -F_val[n_controls:]
        feas = max(0, np.amax(g_val))
        
        # 互补性：lambda * g(u)
        lambda_val = solution[n_controls:]
        comp = np.linalg.norm(g_val * lambda_val, ord=np.inf)

        # print(f"\nPATH求解统计:")
        # print(f"  - 静态性: {stat:.6e}")
        # print(f"  - 可行性: {feas:.6e}")
        # print(f"  - 互补性: {comp:.6e}")

        return {
            'solution': solution,
            'success': success,
            'residual': residual,
            'status': status,
            'stationarity': stat,
            'feasibility': feas,
            'complementarity': comp,
            'iterations': 0  # PATH不提供迭代次数
        }

    def clear_cache(self):
        """
        🆕 清除warm start缓存
        
        用途：
        - 在问题类型变化时重置求解器
        - 在训练epoch之间清理状态
        - 避免缓存累积导致的性能下降
        """
        self.last_solution = None
        self.last_problem_size = None
        # print("  🧹 PATH求解器缓存已清除")


class DifferentiablePATHSolver(torch.autograd.Function):
    """
    可微分的PATH求解器

    实现隐式微分，使得PATH求解器可以在PyTorch计算图中使用
    """

    @staticmethod
    def forward(ctx, F_func, J_func, z0_tensor, lb_tensor=None, ub_tensor=None, solver_kwargs=None):
        """
        前向传播：求解MCP

        Args:
            ctx: autograd上下文
            F_func: 函数F（支持tensor输入）
            J_func: 雅可比矩阵函数J
            z0_tensor: 初始点张量
            lb_tensor: 下界张量
            ub_tensor: 上界张量
            solver_kwargs: 求解器参数

        Returns:
            解张量
        """
        # 转换为numpy
        z0 = z0_tensor.detach().cpu().numpy()
        lb = lb_tensor.detach().cpu().numpy() if lb_tensor is not None else None
        ub = ub_tensor.detach().cpu().numpy() if ub_tensor is not None else None

        solver_kwargs = solver_kwargs or {}

        # 创建numpy版本的函数
        def F_numpy(z):
            z_tensor = torch.from_numpy(z).float().to(z0_tensor.device)
            with torch.no_grad():
                result = F_func(z_tensor)
            return result.detach().cpu().numpy()

        def J_numpy(z):
            z_tensor = torch.from_numpy(z).float().to(z0_tensor.device)
            z_tensor.requires_grad_(True)
            F_val = F_func(z_tensor)
            J_val = torch.autograd.functional.jacobian(
                lambda x: F_func(x), z_tensor
            )
            return J_val.detach().cpu().numpy()

        # 求解
        solver = PATHMCPWrapper(**solver_kwargs)
        result = solver.solve(F_numpy, J_numpy, z0, lb, ub)

        # 转换回tensor
        solution_tensor = torch.from_numpy(result['solution']).float().to(z0_tensor.device)

        # 保存用于反向传播
        ctx.save_for_backward(z0_tensor, solution_tensor)
        ctx.F_func = F_func
        ctx.J_func = J_func
        ctx.solver_kwargs = solver_kwargs

        return solution_tensor

    @staticmethod
    def backward(ctx, grad_output):
        """
        反向传播：隐式微分

        使用隐式函数定理计算梯度
        """
        z0_tensor, solution_tensor = ctx.saved_tensors
        F_func = ctx.F_func

        # 计算雅可比矩阵 ∂F/∂z 在解处
        solution_tensor.requires_grad_(True)
        F_at_sol = F_func(solution_tensor)

        # 计算 ∂F/∂z
        dF_dz = torch.autograd.functional.jacobian(
            lambda z: F_func(z), solution_tensor
        )

        # 使用隐式函数定理：dz/dθ = -(∂F/∂z)^(-1) * (∂F/∂θ)
        # 这里简化处理，假设只对初始点求梯度
        try:
            # 求解线性系统 dF_dz * grad_z = grad_output
            grad_z = torch.linalg.solve(dF_dz.T, grad_output)
        except:
            # 如果矩阵奇异，使用伪逆
            grad_z = torch.linalg.pinv(dF_dz.T) @ grad_output

        return None, None, grad_z, None, None, None


def solve_mcp_differentiable(
    F_func: Callable,
    z0: torch.Tensor,
    lb: Optional[torch.Tensor] = None,
    ub: Optional[torch.Tensor] = None,
    **solver_kwargs
) -> torch.Tensor:
    """
    可微分MCP求解的便捷函数

    Args:
        F_func: 函数F，接受tensor输入返回tensor
        z0: 初始点张量
        lb: 下界张量
        ub: 上界张量
        **solver_kwargs: 求解器参数

    Returns:
        解张量
    """
    # 雅可比函数（在DifferentiablePATHSolver中计算）
    def J_func(z):
        return torch.autograd.functional.jacobian(F_func, z)

    return DifferentiablePATHSolver.apply(
        F_func, J_func, z0, lb, ub, solver_kwargs
    )
